import numpy as np
import pandas as pd
import os
import sys
import pickle
import json
import time
from collections import defaultdict
from pytz import timezone
from datetime import datetime
from sys import getsizeof as getsize
import networkx as nx


def get_span_tag_value(span, key):
    for tag in span['tags']:
        if tag['key'] == key:
            return tag['value']
    return None


def get_endpoint(span):
    http_url = get_span_tag_value(span, 'http.method')
    http_method = get_span_tag_value(span, 'http.url')
    if http_url is None or http_method is None:
        return None
    return http_method + " " + http_url


def is_client(span):
    v = get_span_tag_value(span, 'span.kind')
    if v == 'client':
        return True
    return False


def is_server(span):
    v = get_span_tag_value(span, 'span.kind')
    if v == 'server':
        return True
    return False


def get_root_span(trace):
    for span in trace['spans']:
        if len(span['references']) == 0:
            return span
    return None


def get_proc_svc_map(trace):
    processes = trace['processes']
    proc_svc_map = {}
    for p, v in trace['processes'].items():
        proc_svc_map[p] = v['serviceName']
    return proc_svc_map


def get_node_name(span, service_name):
    if is_client(span):
        return None
    operation = span['operationName']
    return service_name + ":" + operation


def get_node_list(trace):
    node_list = []
    proc_svc_map = get_proc_svc_map(trace)
    for span in trace['spans']:
        node = get_node_name(span, proc_svc_map[span['processID']])
        if node is not None:
            node_list.append(node)
    return sorted(node_list)


def get_trace_metainfo(trace):
    
    valid = True
    trace_id = trace['traceID']
    num_span = len(trace['spans'])
    num_service = len(trace['processes'])
    
    process_map = get_proc_svc_map(trace)
    root_span = get_root_span(trace)
    if root_span is None:
        return {
            'trace_id': trace_id,
            'num_span': num_span,
            'num_service': num_service,
            # 'processes': processes,
            'valid': False,
        }
    
    root_service = None
    root_flags = None
    start_time = None
    duration = None
    root_status_code = None
    
    if root_span is not None:
        root_p = root_span['processID']
        root_service = process_map[root_p]
        root_flags = root_span['flags']
        start_time = root_span['startTime']
        duration = root_span['duration']
        root_status_code = get_span_tag_value(root_span, 'http.status_code')
        if root_status_code is None:
            valid = False
    else:
        valid = False
    
    endpoint = get_endpoint(root_span)
    if endpoint is None:
        valid = False
    
    start_time_str = datetime.fromtimestamp(start_time/1000000).strftime("%Y-%m-%d %H:%M:%S")
    duration = duration / 1000
    
    return {
        'trace_id': trace_id,
        'num_span': num_span,
        'num_service': num_service,
        # 'processes': processes,
        'root_service': root_service,
        'root_flags': root_flags,
        'start_time': start_time_str,
        'duration': str(duration) + " ms",
        'root_status_code': root_status_code,
        'endpoint': endpoint,
        'valid': valid,
    }


def get_span_map(trace):
    proc_svc_map = get_proc_svc_map(trace)
    span_map = {}
    for span in trace['spans']:
        span_id = span['spanID']
        if span['references'] is None or len(span['references']) == 0:
            pid = None
        else:
            pid = span['references'][0]['spanID']
        start_time = span['startTime']
        duration = span['duration']
        status_code = get_span_tag_value(span, 'http.status_code')
        node_name = get_node_name(span, proc_svc_map[span['processID']]) # client span -> None
        span_map[span_id] = {
            'pid': pid,
            'start_time': start_time,
            'duration': duration,
            'status': 0 if status_code == 200 else 1,
            'node_name': node_name,
        }
    return span_map


def construct_graph(trace, directed=False):
    node_list = get_node_list(trace)
    if node_list is None or len(node_list) == 0:
        return False, "invalid node_list"
    
    node_idx_map = {}
    idx_node_map = {}
    for idx, node in enumerate(node_list):
        node_idx_map[node] = idx
        idx_node_map[idx] = node
    
    span_map = get_span_map(trace)
    if span_map is None or len(span_map) == 0 or len(span_map) != len(trace['spans']):
        return False, "invalid span_map"
    
    downs_list_map = {}
    for span_id, span in span_map.items():
        node_name = span['node_name']
        pid = span['pid']
        if pid is not None and pid not in downs_list_map:
            downs_list_map[pid] = []
        if span_id not in downs_list_map:
            downs_list_map[span_id] = []
        if pid is not None:
            downs_list_map[pid].append(span_id)
    
    for span_id, span in span_map.items():
        pid = span['pid']
        if span['node_name'] is not None:  # server span must have an upstream client span
            if pid is None:  # root span
                continue
            if pid not in span_map:
                return False, "The upstream client span {} of server span {} disappear".format(pid, span_id)
        else:  # client span must have an upstream and a downstream server span
            if pid not in span_map:
                return False, "The upstream server span {} of client span {} disappear".format(pid, span_id)
            if span_id not in downs_list_map or len(downs_list_map[span_id]) == 0:
                return False, "{} has no downstream server span".format(span_id)
            if len(downs_list_map[span_id]) > 1:
                return False, "{} has more than one downstream server spans".format(span_id)
    
    # [total_time, rpc_time, local_time, net_time, status, ret_status, out_degree]
    node_feats = np.zeros((len(node_list), 7))
    edge_list = []
    
    for span_id, span in span_map.items():
        node_name = span['node_name']
        if node_name is None:  # client span, not server span
            continue
        
        idx = node_idx_map[node_name]
        pid = span['pid']
        duration = span['duration']
        status = span['status']
        
        # construct features
        duration_up = 0
        if pid is not None:  # not root span
            upper_client_span = span_map[pid]
            duration_up = upper_client_span['duration']
            
            # construct edges
            client_pid = upper_client_span['pid']
            upper_server_span = span_map[client_pid]
            upper_server_node_name = upper_server_span['node_name']
            if upper_server_node_name is None:
                return False, "{} is a client span, not a server span".format(client_pid)
            upper_idx = node_idx_map[upper_server_node_name]
            
            edge_list.append([upper_idx, idx])
            if not directed:
                edge_list.append([idx, upper_idx])
        
        down_client_ids = downs_list_map[span_id]
        down_duration_sum = 0
        down_status_sum = 0
        for child_span_id in down_client_ids:
            child_span = span_map[child_span_id]
            down_duration_sum += child_span['duration']
            down_status_sum += child_span['status']
        
        total_time = duration
        rpc_time = down_duration_sum
        local_time = total_time - rpc_time
        net_time = max(0, duration_up - duration)
        status = 1 if down_status_sum > 0 else 0
        ret_status = status
        out_degree = len(down_client_ids)

        node_feats[idx] = [total_time, rpc_time, local_time, net_time, status, ret_status, out_degree]
        
    return True, (node_feats, edge_list, idx_node_map)


def get_diff_node_pairs(edge, target_edge):
    diff_pairs = {}
    for i in range(len(edge)):
        e1 = edge[i]
        e2 = target_edge[i]
        if e1[0] != e2[0]:
            diff_pairs[min(e1[0], e2[0])] = max(e1[0], e2[0])
        if e1[1] != e2[1]:
            diff_pairs[min(e1[1], e2[1])] = max(e1[1], e2[1])
    return diff_pairs


def repermute_feature(feature, diff_pairs):
    for p1, p2 in diff_pairs.items():
        feature[[p1,p2],:] = feature[[p2,p1],:]